/*********************************************************
          File Name:su.cpp
          Author: Abby Cin
          Mail: abbytsing@gmail.com
          Created Time: Thu 05 Nov 2015 08:35:59 PM CST
**********************************************************/

#include <iostream>
#include <string>
#include <vector>
#include <fstream>
#include <sys/wait.h>
#include <unistd.h>
#include <signal.h>
#include <string.h>
#include <stdlib.h>

#define passwd_path  "/etc/passwd"
pid_t pid = 0;

void handler(int sig)
{
  printf("Received `%s`, killing shell (pid:%d)\n", strsignal(sig), getpid());
  if(kill(pid, SIGKILL) == -1)
    perror("kill");
}

void root()
{
  std::ifstream in(passwd_path);
  std::string line, shell, path;

  if(in.good())
  {
    while(std::getline(in, line))
    {
      if(line.find("root") != line.npos)
        break;
    }
  }
  else
  {
    std::cerr << "Get Default Shell Failed\n";
    exit(1);
  }

  in.close();

  shell = line.substr(line.find_last_of(":") + 1, line.length());
  setenv("HOME", "/root", 1);
  setenv("SHELL", shell.c_str(), 1);
  path = getenv("PATH");
  if(path.find("/usr/sbin") == path.npos)
  {
    if(path.find("/sbin") == path.npos)
      path += ":/sbin";
    path += ":/usr/sbin";
    char p[] = "PATH=''";
    putenv(p);
    setenv("PATH", path.c_str(), 1);
  }

  if(execlp(shell.c_str(), shell.c_str(), NULL) == -1)
  {
    perror("Error");
    exit(1);
  }
}

int main()
{
  if(setuid(0) == -1)
  {
    perror("setuid");
    exit(1);
  }
  if(setgid(0) == -1)
  {
    perror("setgid");
    exit(1);
  }

  switch((pid = fork()))
  {
    case -1:
      std::cerr << "Resource Unavailable\n";
      exit(1);
    case 0:
      root();
    default:
      break;
  }

  sigset_t ex_set;
  sigfillset(&ex_set);
  sigdelset(&ex_set, SIGTERM);
  sigprocmask(SIG_BLOCK, &ex_set, nullptr);
  struct sigaction act;
  sigemptyset(&act.sa_mask);
  sigaddset(&act.sa_mask, SIGTERM);
  act.sa_flags = 0;
  act.sa_handler = handler;
  sigaction(SIGTERM, &act, nullptr);

  wait(NULL);
  return 0;
}
